#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
@version: ??
@author: happiness
@license: Apache Licence 
@contact: happiness_ws@163.com
@site: 
@software: PyCharm
@file: LeNet_5.py
@time: 2017/11/8 10:21

使用leNet-5识别MNIST数据集,它是七层网络结构
"""

import tensorflow as tf
import tensorflow.examples.tutorials.mnist.input_data as input_data
import numpy as np

INPUT_SZIE = 28 * 28
OUTPUT_SIZE = 10

CONV1_DEEP = 32
CONV1_SIZE = 5

CONV2_DEEP = 64
CONV2_SIZE = 5

IMAGE_SIZE = 28
IMAGE_CHANNELS = 1
IMAGE_LABLE = 10

FC_SIZE = 512

BATCH_SIZE = 100

REGULARIZER_RATE = 0.0001
MOVING_AVERAGE_DECAY = 0.99
LEARNING_RATE_BASE = 0.1
DECAY_STEP = 0.5

TRAINING_STEPS = 6000


def inference(input_tensor, trail, regularizer):
    '''
    创建LeNet5卷积神经网络
    :param input_tensor: 输入
    :param trail: 是否是训练数据
    :param regularizer: 正则化
    :return: 返回神经网络
    '''
    with tf.variable_scope("layer1-conv1"):
        conv1_weight = tf.get_variable('weight', shape=[CONV1_SIZE, CONV1_SIZE, IMAGE_CHANNELS, CONV1_DEEP],
                                       dtype=tf.float32,
                                       initializer=tf.truncated_normal_initializer(stddev=0.1))
        conv1_bias = tf.get_variable('bias', shape=[CONV1_DEEP], dtype=tf.float32,
                                     initializer=tf.constant_initializer(0.0))
        conv1 = tf.nn.conv2d(input_tensor, conv1_weight, strides=[1, 1, 1, 1], padding="SAME")
        relu1 = tf.nn.relu(tf.nn.bias_add(conv1, conv1_bias))

    with tf.variable_scope("layer2-pool1"):
        pool1 = tf.nn.max_pool(relu1, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")

    with tf.variable_scope("layer3-conv2"):
        conv2_weight = tf.get_variable('weight', shape=[CONV2_SIZE, CONV2_SIZE, CONV1_DEEP, CONV2_DEEP],
                                       initializer=tf.truncated_normal_initializer(stddev=0.1))
        conv2_bias = tf.get_variable('bias', shape=[CONV2_DEEP], dtype=tf.float32,
                                     initializer=tf.constant_initializer(0.0))

        conv2 = tf.nn.conv2d(pool1, conv2_weight, strides=[1, 1, 1, 1], padding="SAME")
        relu2 = tf.nn.relu(tf.nn.bias_add(conv2, conv2_bias))

    with tf.variable_scope("layer3-pool2"):
        pool2 = tf.nn.max_pool(relu2, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding="SAME")
        pool_shape = pool2.get_shape().as_list()
        nodes = pool_shape[1] * pool_shape[2] * pool_shape[3]
        reshape = tf.reshape(pool2, shape=[pool_shape[0], nodes])

    with tf.variable_scope("layer4-fc1"):
        fc1_wights = tf.get_variable("weight", shape=[nodes, FC_SIZE], dtype=tf.float32,
                                     initializer=tf.truncated_normal_initializer(stddev=0.1))
        fc1_bias = tf.get_variable("bias", shape=[FC_SIZE], dtype=tf.float32, initializer=tf.constant_initializer(0.1))

        if regularizer != None:
            tf.add_to_collection("losses", regularizer(fc1_wights))

        fc1 = tf.nn.relu(tf.matmul(reshape, fc1_wights) + fc1_bias)

        if trail: tf.nn.dropout(fc1, 0.5)

    with tf.variable_scope("layer5-fc2"):
        fc2_weight = tf.get_variable("weight", shape=[FC_SIZE, IMAGE_LABLE], dtype=tf.float32,
                                     initializer=tf.truncated_normal_initializer(stddev=0.1))

        fc2_bias = tf.get_variable("bias", shape=[IMAGE_LABLE], dtype=tf.float32,
                                   initializer=tf.constant_initializer(0.1))

        if regularizer != None:
            tf.add_to_collection("losses", regularizer(fc2_weight))

        logit = tf.matmul(fc1, fc2_weight) + fc2_bias

    return logit


def taril(mnist):
    # 创建输入输出
    x = tf.placeholder(dtype=tf.float32, shape=[BATCH_SIZE, IMAGE_SIZE, IMAGE_SIZE, IMAGE_CHANNELS], name="x-input")
    y_ = tf.placeholder(dtype=tf.float32, shape=[None, IMAGE_LABLE], name="y-input")

    # 设置正则项
    regularizer = tf.contrib.layers.l2_regularizer(REGULARIZER_RATE)
    y = inference(x, False, regularizer)

    # 设置移动平均,对训练的参数做移动平均
    global_step = tf.Variable(0, trainable=False)
    variable_average = tf.train.ExponentialMovingAverage(MOVING_AVERAGE_DECAY, global_step)
    variable_average_op = variable_average.apply(tf.trainable_variables())

    # 定义损失函数
    cross_entroy = tf.nn.sparse_softmax_cross_entropy_with_logits(logits=y, labels=tf.argmax(y_, 1))
    cross_entroy_mean = tf.reduce_mean(cross_entroy)
    loss = cross_entroy_mean + tf.add_n(tf.get_collection("losses"))

    # 设置学习率
    learning_rate = tf.train.exponential_decay(LEARNING_RATE_BASE, global_step=global_step,
                                               decay_steps=mnist.train.num_examples / BATCH_SIZE, decay_rate=DECAY_STEP,
                                               staircase=True)

    trail_step = tf.train.GradientDescentOptimizer(learning_rate=learning_rate).minimize(loss=loss,
                                                                                         global_step=global_step)
    '''
    训练时要执行多个操作时使用control_dependencies或者group
    eg.     with tf.control_dependencies([a, b]):
                c= tf.no_op(name='train')#tf.no_op；什么也不做
                sess.run(c)
    '''
    with tf.control_dependencies([trail_step,variable_average_op ]):
        trail_op = tf.no_op(name="trail")#什么也不做

    with tf.Session() as sess:
        tf.global_variables_initializer().run()

        for i in range(TRAINING_STEPS):
            xs, ys = mnist.train.next_batch(BATCH_SIZE)

            reshaped_xs = np.reshape(xs, (
                BATCH_SIZE,
                IMAGE_SIZE,
                IMAGE_SIZE,
                IMAGE_CHANNELS))
            _, loss_value, step = sess.run([trail_op, loss, global_step], feed_dict={x: reshaped_xs, y_: ys})

            if i % 1000 == 0:
                print("After %d training step(s), loss on training batch is %g." % (step, loss_value))


if __name__ == '__main__':
    mnist = input_data.read_data_sets("./MNIST_DATA/", one_hot=True)
    taril(mnist)
